package ace.module.oauth2.server.core.impl.authentication.ace;

import ace.cmp.core.model.R;
import ace.cmp.security.core.impl.constant.Oauth2Constant;
import ace.module.oauth2.server.api.model.dto.ext.Oauth2AceAuthResultDetailDto;
import ace.module.oauth2.server.core.impl.authentication.ace.handler.Oauth2AceAuthenticationHandler;
import ace.module.oauth2.server.core.impl.authentication.ace.service.AceAuthenticationService;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import lombok.AllArgsConstructor;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.security.authentication.AuthenticationProvider;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.oauth2.core.AuthorizationGrantType;
import org.springframework.security.oauth2.core.ClaimAccessor;
import org.springframework.security.oauth2.core.ClientAuthenticationMethod;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2Error;
import org.springframework.security.oauth2.core.OAuth2ErrorCodes;
import org.springframework.security.oauth2.core.OAuth2RefreshToken;
import org.springframework.security.oauth2.core.OAuth2Token;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.core.oidc.OidcScopes;
import org.springframework.security.oauth2.core.oidc.endpoint.OidcParameterNames;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AccessTokenAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContextHolder;
import org.springframework.security.oauth2.server.authorization.token.DefaultOAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

/**
 * ACE 模式 认证 token.
 *
 * @author caspar
 */
@AllArgsConstructor
@Component
public class Oauth2AceAuthenticationProvider implements AuthenticationProvider {

  private static final String ERROR_URI = "https://datatracker.ietf.org/doc/html/rfc6749#section-5.2";
  private static final OAuth2TokenType ID_TOKEN_TOKEN_TYPE = new OAuth2TokenType(OidcParameterNames.ID_TOKEN);
  private final Log log = LogFactory.getLog(getClass());
  private final OAuth2AuthorizationService authorizationService;
  private final OAuth2TokenGenerator<? extends OAuth2Token> tokenGenerator;
  private final AceAuthenticationService aceAuthenticationService;
  private final Oauth2AceAuthenticationHandler oauth2AceAuthenticationHandler;

  @Override
  public Authentication authenticate(Authentication authentication) throws AuthenticationException {
    Oauth2AceAuthenticationToken oauth2AceAuthenticationToken = (Oauth2AceAuthenticationToken) authentication;

    OAuth2ClientAuthenticationToken clientPrincipal = this.getAuthenticatedClientElseThrowInvalidClient(oauth2AceAuthenticationToken);

    RegisteredClient registeredClient = clientPrincipal.getRegisteredClient();

    // 主认证方法
    R<Oauth2AceAuthResultDetailDto> authenticationResult = this.oauth2AceAuthenticationHandler.authentication(oauth2AceAuthenticationToken);

    String accountId = authenticationResult.check().getAccountId();

    Map<String, String> accountAttributes = Optional.ofNullable(authenticationResult)
        .map(R::getData)
        .map(Oauth2AceAuthResultDetailDto::getAttributes)
        .orElse(new HashMap<>());

    Set<String> authorizedScopes = new HashSet<>();
    if (!CollectionUtils.isEmpty(oauth2AceAuthenticationToken.getRequestedScopes())) {
      for (String requestedScope : oauth2AceAuthenticationToken.getRequestedScopes()) {
        if (!registeredClient.getScopes().contains(requestedScope)) {
          throw new OAuth2AuthenticationException(OAuth2ErrorCodes.INVALID_SCOPE);
        }
      }
      authorizedScopes = new LinkedHashSet<>(oauth2AceAuthenticationToken.getRequestedScopes());
    }

    AuthorizationGrantType authorizationGrantType = aceAuthenticationService.getAuthorizationGrantType();

    OAuth2Authorization authorization =
        OAuth2Authorization.withRegisteredClient(registeredClient)
            .principalName(accountId)
            .authorizedScopes(authorizedScopes)
            .attributes(
                attributes ->
                    attributes.put(Oauth2Constant.TOKEN_ATTRIBUTES_KEY, accountAttributes))
            .authorizationGrantType(authorizationGrantType)
            .build();

    // @formatter:off
    DefaultOAuth2TokenContext.Builder tokenContextBuilder =
        DefaultOAuth2TokenContext.builder()
            .registeredClient(registeredClient)
            .principal(oauth2AceAuthenticationToken)
            .authorizationServerContext(AuthorizationServerContextHolder.getContext())
            .authorization(authorization)
            .authorizedScopes(authorization.getAuthorizedScopes())
            .authorizationGrantType(authorizationGrantType)
            .authorizationGrant(oauth2AceAuthenticationToken);
    // @formatter:on

    OAuth2Authorization.Builder authorizationBuilder = OAuth2Authorization.from(authorization);

    // ----- Access token -----
    OAuth2TokenContext tokenContext =
        tokenContextBuilder.tokenType(OAuth2TokenType.ACCESS_TOKEN).build();
    OAuth2Token generatedAccessToken = this.tokenGenerator.generate(tokenContext);
    if (generatedAccessToken == null) {
      OAuth2Error error =
          new OAuth2Error(
              OAuth2ErrorCodes.SERVER_ERROR,
              "The token generator failed to generate the access token.",
              ERROR_URI);
      throw new OAuth2AuthenticationException(error);
    }

    if (this.log.isTraceEnabled()) {
      this.log.trace("Generated access token");
    }

    OAuth2AccessToken accessToken =
        new OAuth2AccessToken(
            OAuth2AccessToken.TokenType.BEARER,
            generatedAccessToken.getTokenValue(),
            generatedAccessToken.getIssuedAt(),
            generatedAccessToken.getExpiresAt(),
            tokenContext.getAuthorizedScopes());
    if (generatedAccessToken instanceof ClaimAccessor) {
      authorizationBuilder.token(
          accessToken,
          (metadata) ->
              metadata.put(
                  OAuth2Authorization.Token.CLAIMS_METADATA_NAME,
                  ((ClaimAccessor) generatedAccessToken).getClaims()));
    } else {
      authorizationBuilder.accessToken(accessToken);
    }

    // ----- Refresh token -----
    OAuth2RefreshToken refreshToken = null;
    if (registeredClient.getAuthorizationGrantTypes().contains(AuthorizationGrantType.REFRESH_TOKEN)
        &&
        // Do not issue refresh token to public client
        !clientPrincipal.getClientAuthenticationMethod().equals(ClientAuthenticationMethod.NONE)) {

      tokenContext = tokenContextBuilder.tokenType(OAuth2TokenType.REFRESH_TOKEN).build();
      OAuth2Token generatedRefreshToken = this.tokenGenerator.generate(tokenContext);
      if (!(generatedRefreshToken instanceof OAuth2RefreshToken)) {
        OAuth2Error error =
            new OAuth2Error(
                OAuth2ErrorCodes.SERVER_ERROR,
                "The token generator failed to generate the refresh token.",
                ERROR_URI);
        throw new OAuth2AuthenticationException(error);
      }

      if (this.log.isTraceEnabled()) {
        this.log.trace("Generated refresh token");
      }

      refreshToken = (OAuth2RefreshToken) generatedRefreshToken;
      authorizationBuilder.refreshToken(refreshToken);
    }

    // ----- ID token -----
    OidcIdToken idToken;
    if (registeredClient.getScopes().contains(OidcScopes.OPENID)) {
      // @formatter:off
      tokenContext =
          tokenContextBuilder
              .tokenType(ID_TOKEN_TOKEN_TYPE)
              .authorization(
                  authorizationBuilder
                      .build()) // ID token customizer may need access to the access token and/or
              // refresh token
              .build();
      // @formatter:on
      OAuth2Token generatedIdToken = this.tokenGenerator.generate(tokenContext);
      if (!(generatedIdToken instanceof Jwt)) {
        OAuth2Error error =
            new OAuth2Error(
                OAuth2ErrorCodes.SERVER_ERROR,
                "The token generator failed to generate the ID token.",
                ERROR_URI);
        throw new OAuth2AuthenticationException(error);
      }

      if (this.log.isTraceEnabled()) {
        this.log.trace("Generated id token");
      }

      idToken =
          new OidcIdToken(
              generatedIdToken.getTokenValue(),
              generatedIdToken.getIssuedAt(),
              generatedIdToken.getExpiresAt(),
              ((Jwt) generatedIdToken).getClaims());

      authorizationBuilder.token(
          idToken,
          (metadata) ->
              metadata.put(OAuth2Authorization.Token.CLAIMS_METADATA_NAME, idToken.getClaims()));
    } else {
      idToken = null;
    }

    authorization = authorizationBuilder.build();

    this.authorizationService.save(authorization);

    if (this.log.isTraceEnabled()) {
      this.log.trace("Saved authorization");
    }

    Map<String, Object> additionalParameters = Collections.emptyMap();
    if (idToken != null) {
      additionalParameters = new HashMap<>();
      additionalParameters.put(OidcParameterNames.ID_TOKEN, idToken.getTokenValue());
    }

    if (this.log.isTraceEnabled()) {
      this.log.trace("Authenticated token input");
    }

    return new OAuth2AccessTokenAuthenticationToken(
        registeredClient, clientPrincipal, accessToken, refreshToken, additionalParameters);
  }


  @Override
  public boolean supports(Class<?> authentication) {
    return Oauth2AceAuthenticationToken.class.isAssignableFrom(authentication);
  }

  private OAuth2ClientAuthenticationToken getAuthenticatedClientElseThrowInvalidClient(Authentication authentication) {
    OAuth2ClientAuthenticationToken clientPrincipal = null;
    if (OAuth2ClientAuthenticationToken.class.isAssignableFrom(authentication.getPrincipal().getClass())) {
      clientPrincipal = (OAuth2ClientAuthenticationToken) authentication.getPrincipal();
    }
    if (clientPrincipal != null && clientPrincipal.isAuthenticated()) {
      return clientPrincipal;
    }
    throw new OAuth2AuthenticationException(OAuth2ErrorCodes.INVALID_CLIENT);
  }
}
